msg_tool\scripts\kirikiri/
scn.rs

1//! Kirikiri Scene File (.scn)
2use super::mdf::Mdf;
3use crate::ext::io::*;
4use crate::ext::mutex::*;
5use crate::ext::psb::*;
6use crate::scripts::base::*;
7use crate::types::*;
8use crate::utils::encoding::*;
9use anyhow::Result;
10use emote_psb::{PsbReader, PsbWriter};
11use fancy_regex::Regex;
12use std::collections::{HashMap, HashSet};
13use std::io::{Read, Seek};
14use std::path::Path;
15use std::sync::{Arc, Mutex};
16
17#[derive(Debug)]
18/// Kirikiri Scene Script Builder
19pub struct ScnScriptBuilder {}
20
21impl ScnScriptBuilder {
22    /// Creates a new instance of `ScnScriptBuilder`
23    pub fn new() -> Self {
24        Self {}
25    }
26}
27
28impl ScriptBuilder for ScnScriptBuilder {
29    fn default_encoding(&self) -> Encoding {
30        Encoding::Utf8
31    }
32
33    fn build_script(
34        &self,
35        buf: Vec<u8>,
36        filename: &str,
37        _encoding: Encoding,
38        _archive_encoding: Encoding,
39        config: &ExtraConfig,
40        _archive: Option<&Box<dyn Script>>,
41    ) -> Result<Box<dyn Script>> {
42        Ok(Box::new(ScnScript::new(
43            MemReader::new(buf),
44            filename,
45            config,
46        )?))
47    }
48
49    fn build_script_from_file(
50        &self,
51        filename: &str,
52        _encoding: Encoding,
53        _archive_encoding: Encoding,
54        config: &ExtraConfig,
55        _archive: Option<&Box<dyn Script>>,
56    ) -> Result<Box<dyn Script>> {
57        if filename == "-" {
58            let data = crate::utils::files::read_file(filename)?;
59            Ok(Box::new(ScnScript::new(
60                MemReader::new(data),
61                filename,
62                config,
63            )?))
64        } else {
65            let f = std::fs::File::open(filename)?;
66            let reader = std::io::BufReader::new(f);
67            Ok(Box::new(ScnScript::new(reader, filename, config)?))
68        }
69    }
70
71    fn build_script_from_reader(
72        &self,
73        reader: Box<dyn ReadSeek>,
74        filename: &str,
75        _encoding: Encoding,
76        _archive_encoding: Encoding,
77        config: &ExtraConfig,
78        _archive: Option<&Box<dyn Script>>,
79    ) -> Result<Box<dyn Script>> {
80        Ok(Box::new(ScnScript::new(reader, filename, config)?))
81    }
82
83    fn extensions(&self) -> &'static [&'static str] {
84        &["scn"]
85    }
86
87    fn script_type(&self) -> &'static ScriptType {
88        &ScriptType::KirikiriScn
89    }
90
91    fn is_this_format(&self, filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
92        if Path::new(filename)
93            .file_name()
94            .map(|name| {
95                name.to_ascii_lowercase()
96                    .to_string_lossy()
97                    .ends_with(".scn")
98            })
99            .unwrap_or(false)
100            && buf_len >= 4
101            && buf.starts_with(b"PSB\0")
102        {
103            return Some(255);
104        }
105        None
106    }
107}
108
109#[derive(Debug)]
110/// Kirikiri Scene Script
111pub struct ScnScript {
112    psb: VirtualPsbFixed,
113    language_index: usize,
114    languages: Option<Arc<Vec<String>>>,
115    export_chat: bool,
116    filename: String,
117    chat_key: Option<Vec<String>>,
118    chat_json: Option<Arc<HashMap<String, HashMap<String, (String, usize)>>>>,
119    custom_yaml: bool,
120    title: bool,
121    chat_multilang: bool,
122    insert_language: bool,
123}
124
125impl ScnScript {
126    /// Creates a new `ScnScript` from the given reader and filename
127    ///
128    /// * `reader` - The reader containing the PSB or MDF data
129    /// * `filename` - The name of the file (used for error reporting and extension detection)
130    /// * `config` - Extra configuration options
131    pub fn new<R: Read + Seek>(
132        mut reader: R,
133        filename: &str,
134        config: &ExtraConfig,
135    ) -> Result<Self> {
136        let mut header = [0u8; 4];
137        reader.read_exact(&mut header)?;
138        if &header == b"mdf\0" {
139            let mut data = Vec::new();
140            reader.read_to_end(&mut data)?;
141            let decoded = Mdf::unpack(MemReaderRef::new(&data))?;
142            return Self::new(MemReader::new(decoded), filename, config);
143        }
144        reader.rewind()?;
145        let mut psb = PsbReader::open_psb(reader)
146            .map_err(|e| anyhow::anyhow!("Failed to open PSB from {}: {:?}", filename, e))?;
147        let psb = psb
148            .load()
149            .map_err(|e| anyhow::anyhow!("Failed to load PSB from {}: {:?}", filename, e))?;
150        Ok(Self {
151            psb: psb.to_psb_fixed(),
152            language_index: config.kirikiri_language_index.unwrap_or(0),
153            languages: config.kirikiri_languages.clone(),
154            export_chat: config.kirikiri_export_chat,
155            filename: filename.to_string(),
156            chat_key: config.kirikiri_chat_key.clone(),
157            chat_json: config.kirikiri_chat_json.clone(),
158            custom_yaml: config.custom_yaml,
159            title: config.kirikiri_title,
160            chat_multilang: config.kirikiri_chat_multilang,
161            insert_language: config.kirikiri_language_insert,
162        })
163    }
164}
165
166impl Script for ScnScript {
167    fn default_output_script_type(&self) -> OutputScriptType {
168        OutputScriptType::Json
169    }
170
171    fn default_format_type(&self) -> FormatOptions {
172        FormatOptions::None
173    }
174
175    fn is_output_supported(&self, _: OutputScriptType) -> bool {
176        true
177    }
178
179    fn custom_output_extension<'a>(&'a self) -> &'a str {
180        if self.custom_yaml { "yaml" } else { "json" }
181    }
182
183    fn extract_messages(&self) -> Result<Vec<Message>> {
184        let mut messages = Vec::new();
185        let root = self.psb.root();
186        let scenes = root
187            .get_value("scenes")
188            .ok_or(anyhow::anyhow!("scenes not found"))?;
189        let scenes = match scenes {
190            PsbValueFixed::List(list) => list,
191            _ => return Err(anyhow::anyhow!("scenes is not a list")),
192        };
193        let language = if self.language_index != 0 {
194            let index = self.language_index - 1;
195            if let Some(lang) = root["languages"][index].as_str() {
196                Some(lang.to_owned())
197            } else if let Some(languages) = self.languages.as_ref() {
198                if index < languages.len() {
199                    eprintln!(
200                        "WARN: Language code not found in PSB, using from config. Chat messages may not be extracted correctly."
201                    );
202                    crate::COUNTER.inc_warning();
203                    Some(languages[index].to_owned())
204                } else {
205                    None
206                }
207            } else {
208                None
209            }
210        } else {
211            None
212        };
213        if self.language_index != 0 && language.is_none() {
214            eprintln!(
215                "WARN: Language index is set but language code not found in PSB. Chat messages may not be extracted correctly."
216            );
217            crate::COUNTER.inc_warning();
218        }
219        let mut comu = if self.export_chat {
220            Some(ExportMes::new(
221                self.chat_key
222                    .clone()
223                    .unwrap_or(vec!["comumode".to_string()]),
224                if self.chat_multilang {
225                    language.clone()
226                } else {
227                    None
228                },
229            ))
230        } else {
231            None
232        };
233        for (i, oscene) in scenes.iter().enumerate() {
234            let scene = match oscene {
235                PsbValueFixed::Object(obj) => obj,
236                _ => return Err(anyhow::anyhow!("scene at index {} is not an object", i)),
237            };
238            if self.title {
239                if let Some(title) = scene["title"].as_str() {
240                    messages.push(Message {
241                        name: None,
242                        message: title.to_string(),
243                    });
244                }
245                if scene["title"].is_list() {
246                    if let Some(title) = scene["title"][self.language_index].as_str() {
247                        messages.push(Message {
248                            name: None,
249                            message: title.to_string(),
250                        });
251                    }
252                }
253            }
254            if let Some(PsbValueFixed::List(texts)) = scene.get_value("texts") {
255                for (j, text) in texts.iter().enumerate() {
256                    if let PsbValueFixed::List(text) = text {
257                        let values = text.values();
258                        if values.len() <= 1 {
259                            continue; // Skip if there are not enough values
260                        }
261                        let name = &values[0];
262                        let name = match name {
263                            PsbValueFixed::String(s) => Some(s),
264                            PsbValueFixed::Null => None,
265                            _ => return Err(anyhow::anyhow!("name is not a string or null")),
266                        };
267                        let mut display_name;
268                        let mut message;
269                        if matches!(values[1], PsbValueFixed::List(_)) {
270                            display_name = None;
271                            message = &values[1];
272                        } else {
273                            if values.len() <= 2 {
274                                continue; // Skip if there is no message
275                            }
276                            display_name = match &values[1] {
277                                PsbValueFixed::String(s) => Some(s),
278                                PsbValueFixed::Null => None,
279                                _ => {
280                                    return Err(anyhow::anyhow!(
281                                        "display name is not a string or null at {i},{j}"
282                                    ));
283                                }
284                            };
285                            message = &values[2];
286                        }
287                        if matches!(message, PsbValueFixed::List(_)) {
288                            let tmp = message;
289                            if let PsbValueFixed::List(list) = tmp {
290                                if list.len() > self.language_index {
291                                    if let PsbValueFixed::List(data) =
292                                        &list.values()[self.language_index]
293                                    {
294                                        if data.len() >= 2 {
295                                            let data = data.values();
296                                            display_name = match &data[0] {
297                                                PsbValueFixed::String(s) => Some(s),
298                                                PsbValueFixed::Null => None,
299                                                _ => {
300                                                    return Err(anyhow::anyhow!(
301                                                        "display name is not a string or null at {i},{j}"
302                                                    ));
303                                                }
304                                            };
305                                            message = &data[1];
306                                        }
307                                    }
308                                }
309                            }
310                        }
311                        if let PsbValueFixed::String(message) = message {
312                            match name {
313                                Some(name) => {
314                                    let name = match display_name {
315                                        Some(name) => name.string(),
316                                        None => name.string(),
317                                    };
318                                    let message = message.string();
319                                    messages.push(Message {
320                                        name: Some(name.to_string()),
321                                        message: message.replace("\\n", "\n"),
322                                    });
323                                }
324                                None => {
325                                    let message = message.string();
326                                    messages.push(Message {
327                                        name: None,
328                                        message: message.replace("\\n", "\n"),
329                                    });
330                                }
331                            }
332                        }
333                    }
334                }
335            }
336            if let Some(PsbValueFixed::List(selects)) = scene.get_value("selects") {
337                for select in selects.iter() {
338                    if let PsbValueFixed::Object(select) = select {
339                        let mut text = None;
340                        if let Some(PsbValueFixed::List(language)) = select.get_value("language") {
341                            if language.len() > self.language_index {
342                                let v = &language.values()[self.language_index];
343                                if let PsbValueFixed::Object(v) = v {
344                                    text = match v.get_value("text") {
345                                        Some(PsbValueFixed::String(s)) => Some(s),
346                                        Some(PsbValueFixed::Null) => None,
347                                        None => None,
348                                        _ => {
349                                            return Err(anyhow::anyhow!(
350                                                "select text is not a string or null"
351                                            ));
352                                        }
353                                    }
354                                }
355                            }
356                        }
357                        if text.is_none() {
358                            text = match select.get_value("text") {
359                                Some(PsbValueFixed::String(s)) => Some(s),
360                                Some(PsbValueFixed::Null) => None,
361                                None => None,
362                                _ => {
363                                    return Err(anyhow::anyhow!(
364                                        "select text is not a string or null"
365                                    ));
366                                }
367                            };
368                        }
369                        if let Some(text) = text {
370                            let text = text.string();
371                            messages.push(Message {
372                                name: None,
373                                message: text.replace("\\n", "\n"),
374                            });
375                        }
376                    }
377                }
378            }
379            comu.as_mut().map(|c| c.export(&oscene));
380        }
381        if let Some(comu) = comu {
382            if !comu.messages.is_empty() {
383                let mut pb = std::path::PathBuf::from(&self.filename);
384                let key = self
385                    .chat_key
386                    .clone()
387                    .unwrap_or(vec!["comumode".to_string()])
388                    .join("_");
389                let filename = pb
390                    .file_stem()
391                    .map(|s| s.to_string_lossy())
392                    .unwrap_or(std::borrow::Cow::from(&key));
393                pb.set_file_name(format!("{}_{}.json", filename, key));
394                match std::fs::File::create(&pb) {
395                    Ok(mut f) => {
396                        let messages: Vec<String> = comu.messages.into_iter().collect();
397                        if let Err(e) = serde_json::to_writer_pretty(&mut f, &messages) {
398                            eprintln!("Failed to write chat messages to {}: {:?}", pb.display(), e);
399                            crate::COUNTER.inc_warning();
400                        }
401                    }
402                    Err(e) => {
403                        eprintln!(
404                            "Failed to create chat messages file {}: {:?}",
405                            pb.display(),
406                            e
407                        );
408                        crate::COUNTER.inc_warning();
409                    }
410                }
411            }
412        }
413        Ok(messages)
414    }
415
416    fn import_messages<'a>(
417        &'a self,
418        messages: Vec<Message>,
419        file: Box<dyn WriteSeek + 'a>,
420        _filename: &str,
421        _encoding: Encoding,
422        replacement: Option<&'a ReplacementTable>,
423    ) -> Result<()> {
424        let mut mes = messages.iter();
425        let mut cur_mes = mes.next();
426        let mut psb = self.psb.clone();
427        let root = psb.root_mut();
428        if let Some(lang) = &self.languages {
429            let lang = (**lang).clone();
430            root["languages"] = PsbValueFixed::List(PsbListFixed {
431                values: lang
432                    .into_iter()
433                    .map(|s| PsbValueFixed::String(s.into()))
434                    .collect(),
435            });
436        }
437        let language = if self.language_index != 0 {
438            let index = self.language_index - 1;
439            if let Some(lang) = root["languages"][index].as_str() {
440                Some(lang.to_owned())
441            } else {
442                eprintln!(
443                    "WARN: language code not found in PSB. Some functions may not work correctly. Use --kirikiri-languages to specify language codes."
444                );
445                crate::COUNTER.inc_warning();
446                None
447            }
448        } else {
449            None
450        };
451        let ori_lang = if self.insert_language && self.language_index == 0 {
452            if let Some(lang) = root["languages"][0].as_str() {
453                Some(lang.to_owned())
454            } else {
455                None
456            }
457        } else {
458            None
459        };
460        let scenes = &mut root["scenes"];
461        if !scenes.is_list() {
462            return Err(anyhow::anyhow!("scenes is not an array"));
463        }
464        let comu = self.chat_json.as_ref().map(|json| {
465            ImportMes::new(
466                json,
467                replacement,
468                self.chat_key
469                    .clone()
470                    .unwrap_or(vec!["comumode".to_string()]),
471                if self.chat_multilang {
472                    language.clone()
473                } else {
474                    None
475                },
476                self.filename.clone(),
477                if self.chat_multilang {
478                    ori_lang.clone()
479                } else {
480                    None
481                },
482            )
483        });
484        for (i, scene) in scenes.members_mut().enumerate() {
485            if !scene.is_object() {
486                return Err(anyhow::anyhow!("scene at {} is not an object", i));
487            }
488            if self.title {
489                if scene["title"].is_string() {
490                    let m = match cur_mes {
491                        Some(m) => m,
492                        None => {
493                            return Err(anyhow::anyhow!(
494                                "No enough messages. (title at scene {i})"
495                            ));
496                        }
497                    };
498                    let mut title = m.message.clone();
499                    if let Some(replacement) = replacement {
500                        for (key, value) in replacement.map.iter() {
501                            title = title.replace(key, value);
502                        }
503                    }
504                    if self.language_index == 0 {
505                        if self.insert_language {
506                            let ori_title = scene["title"].as_str().unwrap_or("").to_string();
507                            scene["title"].push_member(title);
508                            scene["title"].push_member(ori_title);
509                        } else {
510                            scene["title"].set_string(title);
511                        }
512                    } else {
513                        let ori_title = scene["title"].as_str().unwrap_or("").to_string();
514                        while scene["title"].len() < self.language_index {
515                            scene["title"].push_member(ori_title.clone());
516                        }
517                        if self.insert_language {
518                            scene["title"].insert_member(self.language_index, title);
519                        } else {
520                            scene["title"].push_member(title);
521                        }
522                    }
523                    cur_mes = mes.next();
524                } else if scene["title"].is_list() {
525                    let m = match cur_mes {
526                        Some(m) => m,
527                        None => {
528                            return Err(anyhow::anyhow!(
529                                "No enough messages. (title at scene {i})"
530                            ));
531                        }
532                    };
533                    let mut title = m.message.clone();
534                    if let Some(replacement) = replacement {
535                        for (key, value) in replacement.map.iter() {
536                            title = title.replace(key, value);
537                        }
538                    }
539                    let ori_title = scene["title"][0].as_str().unwrap_or("").to_string();
540                    if self.insert_language {
541                        while scene["title"].len() < self.language_index {
542                            scene["title"].push_member(ori_title.clone());
543                        }
544                        scene["title"].insert_member(self.language_index, title);
545                    } else {
546                        while scene["title"].len() <= self.language_index {
547                            scene["title"].push_member(ori_title.clone());
548                        }
549                        scene["title"][self.language_index].set_string(title);
550                    }
551                    cur_mes = mes.next();
552                }
553            }
554            if scene["texts"].is_list() {
555                for (j, text) in scene["texts"].members_mut().enumerate() {
556                    if text.is_list() {
557                        if text.len() <= 1 {
558                            continue; // Skip if there are not enough values
559                        }
560                        if cur_mes.is_none() {
561                            cur_mes = mes.next();
562                        }
563                        if !text[0].is_string_or_null() {
564                            return Err(anyhow::anyhow!("name is not a string or null"));
565                        }
566                        let has_name = text[0].is_string();
567                        let has_display_name;
568                        if text[1].is_list() {
569                            if self.insert_language {
570                                let ori = text[1][0].clone();
571                                while text[1].len() < self.language_index {
572                                    text[1].push_member(ori.clone());
573                                }
574                                text[1].insert_member(self.language_index, ori.clone());
575                            } else {
576                                while text[1].len() <= self.language_index {
577                                    text[1][self.language_index] = text[1][0].clone();
578                                }
579                            }
580                            if text[1][self.language_index].is_list()
581                                && text[1][self.language_index].len() >= 2
582                            {
583                                if !text[1][self.language_index][0].is_string_or_null() {
584                                    return Err(anyhow::anyhow!(
585                                        "display name is not a string or null"
586                                    ));
587                                }
588                                if text[1][self.language_index][1].is_string() {
589                                    let m = match cur_mes.take() {
590                                        Some(m) => m,
591                                        None => {
592                                            return Err(anyhow::anyhow!(
593                                                "No enough messages. (text {j} at scene {i})"
594                                            ));
595                                        }
596                                    };
597                                    if has_name {
598                                        if let Some(name) = &m.name {
599                                            let mut name = name.clone();
600                                            if let Some(replacement) = replacement {
601                                                for (key, value) in replacement.map.iter() {
602                                                    name = name.replace(key, value);
603                                                }
604                                            }
605                                            text[1][self.language_index][0].set_string(name);
606                                        } else {
607                                            return Err(anyhow::anyhow!(
608                                                "Name is missing for message. (text {j} at scene {i})"
609                                            ));
610                                        }
611                                    }
612                                    let mut message = m.message.clone();
613                                    if let Some(replacement) = replacement {
614                                        for (key, value) in replacement.map.iter() {
615                                            message = message.replace(key, value);
616                                        }
617                                    }
618                                    text[1][self.language_index][1]
619                                        .set_string(message.replace("\n", "\\n"));
620                                    // text length
621                                    text[1][self.language_index][2]
622                                        .set_i64(message.chars().count() as i64);
623                                    text[1][self.language_index][3]
624                                        .set_string(get_save_message(&message, true));
625                                    text[1][self.language_index][4]
626                                        .set_string(get_save_message(&message, false));
627                                }
628                            }
629                        } else {
630                            if text.len() <= 2 {
631                                continue; // Skip if there is no message
632                            }
633                            if !text[1].is_string_or_null() {
634                                return Err(anyhow::anyhow!(
635                                    "display name is not a string or null"
636                                ));
637                            }
638                            has_display_name = text[1].is_string();
639                            if text[2].is_string() {
640                                let m = match cur_mes.take() {
641                                    Some(m) => m,
642                                    None => {
643                                        return Err(anyhow::anyhow!(
644                                            "No enough messages.(text {j} at scene {i})"
645                                        ));
646                                    }
647                                };
648                                if has_name {
649                                    if let Some(name) = &m.name {
650                                        let mut name = name.clone();
651                                        if let Some(replacement) = replacement {
652                                            for (key, value) in replacement.map.iter() {
653                                                name = name.replace(key, value);
654                                            }
655                                        }
656                                        if has_display_name {
657                                            text[1].set_string(name);
658                                        } else {
659                                            text[0].set_string(name);
660                                        }
661                                    } else {
662                                        return Err(anyhow::anyhow!(
663                                            "Name is missing for message.(text {j} at scene {i})"
664                                        ));
665                                    }
666                                }
667                                let mut message = m.message.clone();
668                                if let Some(replacement) = replacement {
669                                    for (key, value) in replacement.map.iter() {
670                                        message = message.replace(key, value);
671                                    }
672                                }
673                                text[2].set_string(message.replace("\n", "\\n"));
674                            } else if text[2].is_list() {
675                                if self.insert_language {
676                                    let ori = text[2][0].clone();
677                                    while text[2].len() < self.language_index {
678                                        text[2].push_member(ori.clone());
679                                    }
680                                    text[2].insert_member(self.language_index, ori.clone());
681                                } else {
682                                    while text[2].len() <= self.language_index {
683                                        text[2][self.language_index] = text[2][0].clone();
684                                    }
685                                }
686                                if text[2][self.language_index].is_list()
687                                    && text[2][self.language_index].len() >= 2
688                                {
689                                    if !text[2][self.language_index][0].is_string_or_null() {
690                                        return Err(anyhow::anyhow!(
691                                            "display name is not a string or null"
692                                        ));
693                                    }
694                                    if text[2][self.language_index][1].is_string() {
695                                        let m = match cur_mes.take() {
696                                            Some(m) => m,
697                                            None => {
698                                                return Err(anyhow::anyhow!(
699                                                    "No enough messages.(text {j} at scene {i})"
700                                                ));
701                                            }
702                                        };
703                                        if has_name {
704                                            if let Some(name) = &m.name {
705                                                let mut name = name.clone();
706                                                if let Some(replacement) = replacement {
707                                                    for (key, value) in replacement.map.iter() {
708                                                        name = name.replace(key, value);
709                                                    }
710                                                }
711                                                text[2][self.language_index][0].set_string(name);
712                                            } else {
713                                                return Err(anyhow::anyhow!(
714                                                    "Name is missing for message.(text {j} at scene {i})"
715                                                ));
716                                            }
717                                        }
718                                        let mut message = m.message.clone();
719                                        if let Some(replacement) = replacement {
720                                            for (key, value) in replacement.map.iter() {
721                                                message = message.replace(key, value);
722                                            }
723                                        }
724                                        text[2][self.language_index][1]
725                                            .set_string(message.replace("\n", "\\n"));
726                                        text[2][self.language_index][2]
727                                            .set_i64(message.chars().count() as i64);
728                                        text[2][self.language_index][3]
729                                            .set_string(get_save_message(&message, true));
730                                        text[2][self.language_index][4]
731                                            .set_string(get_save_message(&message, false));
732                                    }
733                                }
734                            }
735                        }
736                    }
737                }
738            }
739            if scene["selects"].is_list() {
740                for select in scene["selects"].members_mut() {
741                    if select.is_object() {
742                        if cur_mes.is_none() {
743                            cur_mes = mes.next();
744                        }
745                        if self.language_index != 0
746                            && {
747                                if self.insert_language {
748                                    while select["language"].len() < self.language_index {
749                                        // TenShiSouZou
750                                        // first block is null
751                                        if select["language"].len() == 0 {
752                                            select["language"].push_member(PsbValueFixed::Null);
753                                            continue;
754                                        }
755                                        let mut obj = PsbObjectFixed::new();
756                                        obj["text"].set_str("");
757                                        obj["speechtext"].set_str("");
758                                        obj["searchtext"].set_str("");
759                                        obj["textlength"].set_i64(0);
760                                        select["language"][self.language_index].set_obj(obj);
761                                    }
762                                    let mut obj = PsbObjectFixed::new();
763                                    obj["text"].set_str("");
764                                    obj["speechtext"].set_str("");
765                                    obj["searchtext"].set_str("");
766                                    obj["textlength"].set_i64(0);
767                                    select["language"].insert_member(self.language_index, obj);
768                                } else {
769                                    while select["language"].len() <= self.language_index {
770                                        // TenShiSouZou
771                                        // first block is null
772                                        if select["language"].len() == 0 {
773                                            select["language"].push_member(PsbValueFixed::Null);
774                                            continue;
775                                        }
776                                        let mut obj = PsbObjectFixed::new();
777                                        obj["text"].set_str("");
778                                        obj["speechtext"].set_str("");
779                                        obj["searchtext"].set_str("");
780                                        obj["textlength"].set_i64(0);
781                                        select["language"][self.language_index].set_obj(obj);
782                                    }
783                                }
784                                true
785                            }
786                            && select["language"][self.language_index].is_object()
787                        {
788                            let lang_obj = &mut select["language"][self.language_index];
789                            if lang_obj["text"].is_string() {
790                                let m = match cur_mes.take() {
791                                    Some(m) => m,
792                                    None => {
793                                        return Err(anyhow::anyhow!("No enough messages."));
794                                    }
795                                };
796                                let mut text = m.message.clone();
797                                if let Some(replacement) = replacement {
798                                    for (key, value) in replacement.map.iter() {
799                                        text = text.replace(key, value);
800                                    }
801                                }
802                                lang_obj["text"].set_string(text.replace("\n", "\\n"));
803                                lang_obj["speechtext"].set_string(get_save_message(&text, true));
804                                lang_obj["searchtext"].set_string(get_save_message(&text, false));
805                                lang_obj["textlength"].set_i64(text.chars().count() as i64);
806                                continue;
807                            }
808                        } else if select["text"].is_string() {
809                            let m = match cur_mes.take() {
810                                Some(m) => m,
811                                None => {
812                                    return Err(anyhow::anyhow!("No enough messages."));
813                                }
814                            };
815                            let mut text = m.message.clone();
816                            if let Some(replacement) = replacement {
817                                for (key, value) in replacement.map.iter() {
818                                    text = text.replace(key, value);
819                                }
820                            }
821                            if self.insert_language {
822                                let ori_text = select["text"].as_str().unwrap_or("").to_string();
823                                let mut obj = PsbObjectFixed::new();
824                                obj["text"].set_string(ori_text.replace("\n", "\\n"));
825                                obj["speechtext"].set_string(get_save_message(&ori_text, true));
826                                obj["searchtext"].set_string(get_save_message(&ori_text, false));
827                                obj["textlength"].set_i64(ori_text.chars().count() as i64);
828                                if select["language"].len() < 1 {
829                                    select["language"].push_member(PsbValueFixed::Null);
830                                }
831                                select["language"].insert_member(1, obj);
832                            }
833                            select["text"].set_string(text.replace("\n", "\\n"));
834                        }
835                    }
836                }
837            }
838            comu.as_ref().map(|c| c.import(scene));
839        }
840        if cur_mes.is_some() || mes.next().is_some() {
841            return Err(anyhow::anyhow!("Some messages were not processed."));
842        }
843        let psb = psb.to_psb(true);
844        let writer = PsbWriter::new(psb, file);
845        writer.finish().map_err(|e| {
846            anyhow::anyhow!("Failed to write PSB to file {}: {:?}", self.filename, e)
847        })?;
848        Ok(())
849    }
850
851    fn custom_export(&self, filename: &Path, encoding: Encoding) -> Result<()> {
852        let s = if self.custom_yaml {
853            serde_yaml_ng::to_string(&self.psb)
854                .map_err(|e| anyhow::anyhow!("Failed to serialize to YAML: {}", e))?
855        } else {
856            json::stringify_pretty(self.psb.to_json(), 2)
857        };
858        let mut f = crate::utils::files::write_file(filename)?;
859        let b = encode_string(encoding, &s, false)?;
860        f.write_all(&b)?;
861        Ok(())
862    }
863
864    fn custom_import<'a>(
865        &'a self,
866        custom_filename: &'a str,
867        file: Box<dyn WriteSeek + 'a>,
868        _encoding: Encoding,
869        output_encoding: Encoding,
870    ) -> Result<()> {
871        let data = crate::utils::files::read_file(custom_filename)?;
872        let s = decode_to_string(output_encoding, &data, true)?;
873        let psb = if self.custom_yaml {
874            let data: VirtualPsbFixedData = serde_yaml_ng::from_str(&s)
875                .map_err(|e| anyhow::anyhow!("Failed to deserialize YAML: {}", e))?;
876            let mut psb = self.psb.clone();
877            psb.set_data(data);
878            psb.to_psb(true)
879        } else {
880            let json = json::parse(&s)?;
881            let mut psb = self.psb.clone();
882            psb.from_json(&json)?;
883            psb.to_psb(true)
884        };
885        let writer = PsbWriter::new(psb, file);
886        writer.finish().map_err(|e| {
887            anyhow::anyhow!("Failed to write PSB to file {}: {:?}", self.filename, e)
888        })?;
889        Ok(())
890    }
891}
892
893#[derive(Debug)]
894struct ExportMes {
895    pub messages: HashSet<String>,
896    pub key: HashSet<String>,
897    text_key: String,
898}
899
900impl ExportMes {
901    pub fn new(key: Vec<String>, language: Option<String>) -> Self {
902        Self {
903            messages: HashSet::new(),
904            key: HashSet::from_iter(key.into_iter()),
905            text_key: language.map_or_else(|| String::from("text"), |s| format!("text_{}", s)),
906        }
907    }
908
909    pub fn export(&mut self, value: &PsbValueFixed) {
910        match value {
911            PsbValueFixed::Object(obj) => {
912                for (k, v) in obj.iter() {
913                    if self.key.contains(k) {
914                        if let PsbValueFixed::List(list) = v {
915                            for item in list.iter() {
916                                if let PsbValueFixed::Object(obj) = item {
917                                    if let Some(s) = obj[&self.text_key].as_str() {
918                                        self.messages.insert(s.replace("\\n", "\n"));
919                                    } else if let Some(s) = obj["text"].as_str() {
920                                        self.messages.insert(s.replace("\\n", "\n"));
921                                    }
922                                }
923                            }
924                        }
925                    } else {
926                        self.export(v);
927                    }
928                }
929            }
930            PsbValueFixed::List(list) => {
931                let list = list.values();
932                if list.len() > 1 {
933                    if let PsbValueFixed::String(s) = &list[0] {
934                        if self.key.contains(s.string()) {
935                            for i in 1..list.len() {
936                                if let PsbValueFixed::String(s) = &list[i - 1] {
937                                    if s.string() == &self.text_key {
938                                        if let PsbValueFixed::String(text) = &list[i] {
939                                            self.messages
940                                                .insert(text.string().replace("\\n", "\n"));
941                                        }
942                                    }
943                                }
944                            }
945                            if self.text_key == "text" {
946                                return;
947                            }
948                            for i in 1..list.len() {
949                                if let PsbValueFixed::String(s) = &list[i - 1] {
950                                    if s.string() == "text" {
951                                        if let PsbValueFixed::String(text) = &list[i] {
952                                            self.messages
953                                                .insert(text.string().replace("\\n", "\n"));
954                                        }
955                                    }
956                                }
957                            }
958                            return;
959                        }
960                    }
961                }
962                for item in list {
963                    self.export(item);
964                }
965            }
966            _ => {}
967        }
968    }
969}
970
971lazy_static::lazy_static! {
972    static ref DUP_WARN_SHOWN: Mutex<HashSet<(String, usize, String)>> = Mutex::new(HashSet::new());
973    static ref NOT_FOUND_WARN_SHOWN: Mutex<HashSet<String>> = Mutex::new(HashSet::new());
974}
975
976fn warn_dup(original: String, count: usize, filename: String) {
977    let mut guard = DUP_WARN_SHOWN.lock_blocking();
978    if guard.contains(&(original.clone(), count, filename.clone())) {
979        return;
980    }
981    eprintln!(
982        "Warning: chat message '{}' has {} duplicates in translation table '{}'. Using the first one.",
983        original, count, filename
984    );
985    crate::COUNTER.inc_warning();
986    guard.insert((original.clone(), count, filename.clone()));
987}
988
989fn warn_not_found(original: String) {
990    let mut guard = NOT_FOUND_WARN_SHOWN.lock_blocking();
991    if guard.contains(&original) {
992        return;
993    }
994    eprintln!(
995        "Warning: chat message '{}' not found in translation table.",
996        original
997    );
998    crate::COUNTER.inc_warning();
999    guard.insert(original);
1000}
1001
1002#[derive(Debug)]
1003struct ImportMes<'a> {
1004    messages: &'a Arc<HashMap<String, HashMap<String, (String, usize)>>>,
1005    replacement: Option<&'a ReplacementTable>,
1006    key: HashSet<String>,
1007    text_key: String,
1008    filename: String,
1009    ori_text_key: Option<String>,
1010}
1011
1012impl<'a> ImportMes<'a> {
1013    pub fn new(
1014        messages: &'a Arc<HashMap<String, HashMap<String, (String, usize)>>>,
1015        replacement: Option<&'a ReplacementTable>,
1016        key: Vec<String>,
1017        lang: Option<String>,
1018        filename: String,
1019        ori_lang: Option<String>,
1020    ) -> Self {
1021        Self {
1022            messages,
1023            replacement,
1024            key: HashSet::from_iter(key.into_iter()),
1025            text_key: lang.map_or_else(|| String::from("text"), |s| format!("text_{}", s)),
1026            filename: std::path::Path::new(&filename)
1027                .file_stem()
1028                .map(|s| s.to_string_lossy().to_string())
1029                .unwrap_or_else(|| "global".to_string()),
1030            ori_text_key: ori_lang.map(|s| format!("text_{}", s)),
1031        }
1032    }
1033
1034    fn get_message(&self, original: &str) -> Option<String> {
1035        if let Some(global) = self.messages.get(&self.filename) {
1036            if let Some(text) = global.get(original) {
1037                if text.1 > 1 {
1038                    warn_dup(original.to_string(), text.1, self.filename.clone());
1039                }
1040                return Some(text.0.clone());
1041            }
1042        }
1043        if self.filename == "global" {
1044            return None;
1045        }
1046        if let Some(file) = self.messages.get("global") {
1047            if let Some(text) = file.get(original) {
1048                if text.1 > 1 {
1049                    warn_dup(original.to_string(), text.1, "global".to_string());
1050                }
1051                return Some(text.0.clone());
1052            }
1053        }
1054        None
1055    }
1056
1057    pub fn import(&self, value: &mut PsbValueFixed) {
1058        match value {
1059            PsbValueFixed::Object(obj) => {
1060                for (k, v) in obj.iter_mut() {
1061                    if self.key.contains(k) {
1062                        for obj in v.members_mut() {
1063                            if let Some(text) = obj[&self.text_key].as_str() {
1064                                if let Some(replace_text) = self.get_message(text) {
1065                                    let mut text = replace_text.clone();
1066                                    if let Some(replacement) = self.replacement {
1067                                        for (key, value) in replacement.map.iter() {
1068                                            text = text.replace(key, value);
1069                                        }
1070                                    }
1071                                    if self.text_key == "text" {
1072                                        if let Some(ori_key) = &self.ori_text_key {
1073                                            let ori_text =
1074                                                obj["text"].as_str().unwrap_or("").to_string();
1075                                            obj[ori_key].set_string(ori_text);
1076                                        }
1077                                    }
1078                                    obj[&self.text_key].set_string(text.replace("\n", "\\n"));
1079                                    continue;
1080                                } else {
1081                                    warn_not_found(text.to_string());
1082                                }
1083                            }
1084                            if let Some(text) = obj["text"].as_str() {
1085                                if let Some(replace_text) = self.get_message(text) {
1086                                    let mut text = replace_text.clone();
1087                                    if let Some(replacement) = self.replacement {
1088                                        for (key, value) in replacement.map.iter() {
1089                                            text = text.replace(key, value);
1090                                        }
1091                                    }
1092                                    if let Some(ori_key) = &self.ori_text_key {
1093                                        let ori_text =
1094                                            obj["text"].as_str().unwrap_or("").to_string();
1095                                        obj[ori_key].set_string(ori_text);
1096                                    }
1097                                    obj[&self.text_key].set_string(text.replace("\n", "\\n"));
1098                                } else {
1099                                    warn_not_found(text.to_string());
1100                                }
1101                            }
1102                        }
1103                    } else {
1104                        self.import(v);
1105                    }
1106                }
1107            }
1108            PsbValueFixed::List(list) => {
1109                if list.len() > 1 {
1110                    if list[0].as_str().map_or(false, |s| self.key.contains(s)) {
1111                        for i in 1..list.len() {
1112                            if list[i - 1] == self.text_key {
1113                                if let Some(text) = list[i].as_str() {
1114                                    if let Some(replace_text) = self.get_message(text) {
1115                                        let mut text = replace_text.clone();
1116                                        if let Some(replacement) = self.replacement {
1117                                            for (key, value) in replacement.map.iter() {
1118                                                text = text.replace(key, value);
1119                                            }
1120                                        }
1121                                        if self.text_key == "text" {
1122                                            if let Some(ori_key) = &self.ori_text_key {
1123                                                let len = list.len();
1124                                                let ori_text =
1125                                                    list[i].as_str().unwrap_or("").to_string();
1126                                                list[len].set_str(ori_key);
1127                                                list[len + 1].set_string(ori_text);
1128                                            }
1129                                        }
1130                                        list[i].set_string(text.replace("\n", "\\n"));
1131                                        return;
1132                                    } else {
1133                                        warn_not_found(text.to_string());
1134                                    }
1135                                }
1136                            }
1137                        }
1138                        if self.text_key == "text" {
1139                            return;
1140                        }
1141                        for i in 1..list.len() {
1142                            if list[i - 1] == "text" {
1143                                if let Some(text) = list[i].as_str() {
1144                                    if let Some(replace_text) = self.get_message(text) {
1145                                        let mut text = replace_text.clone();
1146                                        if let Some(replacement) = self.replacement {
1147                                            for (key, value) in replacement.map.iter() {
1148                                                text = text.replace(key, value);
1149                                            }
1150                                        }
1151                                        let len = list.len();
1152                                        list[len].set_str(&self.text_key);
1153                                        list[len + 1].set_string(text.replace("\n", "\\n"));
1154                                        return;
1155                                    } else {
1156                                        warn_not_found(text.to_string());
1157                                    }
1158                                }
1159                            }
1160                        }
1161                        return;
1162                    }
1163                }
1164                for item in list.iter_mut() {
1165                    self.import(item);
1166                }
1167            }
1168            _ => {}
1169        }
1170    }
1171}
1172
1173lazy_static::lazy_static! {
1174    static ref CONTROL: Regex = Regex::new("%[^%;]*;").unwrap();
1175    static ref RUBY: Regex = Regex::new(r"(?<!\\)\[([^\]]*)\](.?)").unwrap();
1176    static ref COLOR: Regex = Regex::new(r"#[0-9a-fA-F]{6,8};").unwrap();
1177}
1178
1179fn get_save_message(s: &str, in_ruby: bool) -> String {
1180    let mut s = s.replace("\n", "");
1181    s = CONTROL.replace_all(&s, "").to_string();
1182    s = COLOR.replace_all(&s, "").to_string();
1183    s = RUBY
1184        .replace_all(&s, if in_ruby { "$1" } else { "$2" })
1185        .to_string();
1186    s.replace("%r", "").replace("\\[", "[")
1187}
1188
1189#[test]
1190fn test_get_save_message() {
1191    let s = "%n;Test\n[ruby]测[test\\]试%ok;[ok]";
1192    assert_eq!(get_save_message(s, true), "Testrubytest\\ok");
1193    assert_eq!(get_save_message(s, false), "Test测试");
1194    let another = "[Start]a";
1195    assert_eq!(get_save_message(another, true), "Start");
1196    assert_eq!(get_save_message(another, false), "a");
1197    let escaped = "\\[Start]a";
1198    assert_eq!(get_save_message(escaped, true), "[Start]a");
1199    assert_eq!(get_save_message(escaped, false), "[Start]a");
1200    let real_word = "「こんな、感じとか、ですか……? うっふ~ん……%f$ハート$;#00ffadd6;♥%r」";
1201    assert_eq!(
1202        get_save_message(real_word, true),
1203        "「こんな、感じとか、ですか……? うっふ~ん……♥」"
1204    );
1205    let s = "「あっは%f$ハート$;#00ffadd6;♥%r 凄い出してくれてる%f$ハート$;#00ffadd6;♥%r」";
1206    assert_eq!(
1207        get_save_message(s, true),
1208        "「あっは♥ 凄い出してくれてる♥」"
1209    );
1210}